from pydub import AudioSegment
import webrtcvad
import numpy as np

def read_wav_as_mono_pcm(path, target_sr=16000):
    # 读取任意音频文件，并将其转换为16kHz, 16位, 单声道
    audio = AudioSegment.from_file(path)  # pydub会根据文件自动解码
    audio = audio.set_frame_rate(target_sr).set_channels(1).set_sample_width(2)
    # audio.raw_data 是 bytes 格式，即PCM原始数据
    return audio.raw_data, audio.frame_rate


def frame_generator(frame_duration_ms, audio_data, sample_rate):
    """
    将原始PCM数据切成一帧帧，每帧frame_duration_ms长度
    返回一个生成器，每次yield一个小段的PCM数据
    """
    n_bytes_per_sample = 2  # 16-bit = 2 bytes
    bytes_per_frame = int(sample_rate * (frame_duration_ms / 1000.0) * n_bytes_per_sample)

    # total_len = 总字节数
    total_len = len(audio_data)
    offset = 0
    while offset + bytes_per_frame <= total_len:
        yield audio_data[offset:offset + bytes_per_frame]
        offset += bytes_per_frame

def vad_collector(vad, frames, sample_rate, frame_duration_ms=30):
    """
    使用webrtcvad对每帧进行判断，将连续的语音帧合并输出。
    """
    voiced_frames = []
    for frame in frames:
        # is_speech 返回 True or False
        if vad.is_speech(frame, sample_rate):
            voiced_frames.append(frame)
        else:
            # 如果当前帧判定为静音/噪声，但 voiced_frames 不为空，
            # 说明前面在语音段，现在语音段结束，这时候可以yield。
            if voiced_frames:
                yield b''.join(voiced_frames)
                voiced_frames = []
    # 最后如果结尾时还有语音帧未输出，补一下
    if voiced_frames:
        yield b''.join(voiced_frames)

def run_vad_on_wav(input_path, aggressiveness=3, frame_duration_ms=30):
    """
    使用WebRTC VAD对输入音频进行语音活动检测，
    返回多个语音片段的PCM原始数据列表.
    aggressiveness：VAD的敏感度，取值0~3，值越高对噪声越敏感，过滤越严格
    """
    # 1. 读取音频，并确保是16k, 16位, 单声道 PCM
    raw_audio, sample_rate = read_wav_as_mono_pcm(input_path, 16000)

    # 2. 初始化 VAD
    vad = webrtcvad.Vad()
    vad.set_mode(aggressiveness)
    # mode=0 最宽松，mode=3 最严格，需要根据实际噪声环境调参

    # 3. 按帧切分音频
    frames = frame_generator(frame_duration_ms, raw_audio, sample_rate)

    # 4. 收集语音帧
    segments = []
    for voiced_chunk in vad_collector(vad, frames, sample_rate, frame_duration_ms):
        segments.append(voiced_chunk)
    return segments

def pcm_to_ndarray(raw_pcm, sample_rate=16000):
    """
    将原始PCM (16-bit, mono, sample_rate) 转为numpy数组
    """
    audio_segment = AudioSegment(
        raw_pcm,
        frame_rate=sample_rate,
        sample_width=2,  # 16-bit
        channels=1
    )
    samples = audio_segment.get_array_of_samples()
    return np.array(samples).astype(np.float32) / 32768.0  # 归一化到[-1,1]

